from enum import Enum
from typing import Optional, List, Dict, Any
from pydantic import BaseModel
from app.schemas.metrics_schema import Chart

# 任务状态定义
class TaskStatus(str, Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING" 
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"

# 风险等级定义
class RiskLevel(str, Enum):
    LOW = "LOW"
    MEDIUM = "MEDIUM"
    HIGH = "HIGH"
    CRITICAL = "CRITICAL"

# DRPA算法输入参数
class DRPAInputParams(BaseModel):
    data_flows: List[Dict[str, Any]]  # 数据流列表
    time_window: int = 60  # 时间窗口(秒)
    risk_threshold: float = 0.7  # 风险阈值
    sensitivity_weights: Dict[str, float] = {  # 敏感度权重
        "data_type": 0.3,
        "access_frequency": 0.2, 
        "user_privilege": 0.3,
        "historical_incidents": 0.2
    }

# 时间窗口的评估结果
class WindowAssessment(BaseModel):
    window_id: int
    start_time: int
    end_time: int
    risk_score: float
    risk_level: RiskLevel
    flagged_flows: List[str]
    mitigation_actions: List[str]

# DRPA算法输出参数
class DRPAOutputParams(BaseModel):
    assessments: List[WindowAssessment]
    overall_risk_score: float
    high_risk_patterns: List[Dict[str, Any]]
    recommended_actions: List[str]

# 输入参数封装
class InputParams(BaseModel):
    drpa_params: Optional[DRPAInputParams] = None

# 输出参数封装
class OutputParams(BaseModel):
    drpa_results: Optional[DRPAOutputParams] = None

# 算法请求
class AlgorithmRequest(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    input_params: InputParams

# 算法中间响应
class AlgorithmMiddleResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    input_params: InputParams
    error_message: Optional[str] = None
    metrics: List[Chart] = []

# 算法最终响应
class AlgorithmResponse(BaseModel):
    task_id: str
    task_callback_url: Optional[str] = None
    task_status: TaskStatus
    task_progress: int = 0
    task_logs: Optional[str] = None
    error_message: Optional[str] = None
    output_params: OutputParams
    metrics: List[Chart] = [] 